print_results <- function(data) {
  outcomes <- data$outcome
  ests <- data$pretty_label
  ss <- data$std_error
  ds <- data$pretty_cohens_d
  pvals <- data$pretty_fdr_p
  out <- sapply(1:6, function(i) paste0(
    outcomes[i], ": ", ests[i],
    " (stderr = " , sprintf("%0i", round(ss[i], 0)),
    ", cohensd = ", ds[i],
    ", pfdr = ", pvals[i],
    "); "))
  out <- gsub("-", "\u2212", out)
  cat(out)
}
print_het_results <- function(i, models, var) {
  x <- models[[i]]
  y <- x$coef[var, c(1, 2, 4)]
  out <- paste0(
    outcomes[i], ": ", sprintf("%0i", round(y[1], 0)),
    " (stderr = " , sprintf("%0i", round(y[2], 0)),
    ", uncorrectedp = ", sprintf("%.04f", round(y[3], 4)),
    "); ")
  out <- gsub("-", "\u2212", out)
  cat(out)
}
print_med_results <- function(i, models) {
  x <- summary(models[[i]])
  y <- c(x$d1, x$d1.ci, x$d1.p)
  out <- paste0(outcomes[i], ": ", sprintf("%.01f", round(y[1], 1)),
    " (CI = [" , sprintf("%.01f", round(y[2], 1)),
    ", " , sprintf("%.01f", round(y[3], 1)),
    "], uncorrectedp = ", sprintf("%.04f", round(y[4], 4)),
    "); ")
  out <- gsub("-", "\u2212", out)
  cat(out)
}


weightassess <- function(x, y, z) {
  w <- anesrake::weightassess(x, y, z)
  max(abs(unlist(sapply(w, function(ww) head(ww, -1)[, 7]))))
}
rake <- function(x, y, z) {
  disc <- 1e2
  cap <- 1
  while (disc > .01) {
    cap <- cap + 1
    cat(cap, "\n")
    rake_fit <- anesrake::anesrake(x, y, caseid = z, cap = cap,
      maxit = 1e4)
    disc <- weightassess(x, y, rake_fit$weightvec)
  }
  rake_fit
}

add_region <- function(data) {
  data[state %in% c("Connecticut", "Maine", "Rhode Island", "Vermont",
    "New Hampshire", "Massachusetts"),
    region := "New England"]
  data[state %in% c("Colorado", "New Mexico", "Montana", "Arizona",
    "Utah", "Nevada", "Wyoming", "Idaho"),
    region := "Mountain"]
  data[state %in% c("Minnesota", "Missouri", "Kansas", "Iowa", "Nebraska",
    "North Dakota", "South Dakota"),
    region := "West North Central"]
  data[state %in% c("Illinois", "Ohio", "Indiana", "Michigan", "Wisconsin"),
    region := "East North Central"]
  data[state %in% c("Oklahoma", "Texas", "Arkansas", "Louisiana"),
    region := "West South Central"]
  data[state %in% c("Kentucky", "Tennessee", "Alabama", "Mississippi"),
    region := "East South Central"]
  data[state %in% c("New York", "New Jersey", "Pennsylvania"),
    region := "Middle Atlantic"]
  data[state %in% c("Georgia", "Florida", "North Carolina", "South Carolina",
    "Delaware", "Maryland", "District of Columbia", "Virginia",
    "West Virginia"),
    region := "South Atlantic"]
  data[state %in% c("California", "Washington", "Oregon", "Hawaii", "Alaska"),
    region := "Pacific"]
  data[, division := region]
  data[division %in% c("New England", "Middle Atlantic"), region := "Northeast"]
  data[division %in% c("East North Central", "West North Central"),
    region := "Midwest"]
  data[division %in% c("South Atlantic", "West South Central",
    "East South Central"),
    region := "South"]
  data[division %in% c("Pacific", "Mountain"), region := "West"]
  invisible()
}

make_targets <- function(data) {
  categories <- data[,
    .(weight = sum(commonweight)),
    .(
      age_cat = factor(
        ifelse((2022 - birthyr) %in% 18:29, "18-29",
        ifelse((2022 - birthyr) %in% 30:44, "30-44",
        ifelse((2022 - birthyr) %in% 45:64, "45-64",
        ifelse((2022 - birthyr) %in% 65:200, "65+", "missing"))))),
      # age_cat = factor(
      #   ifelse((2022 - birthyr) %in% 18:24, "18-24",
      #   ifelse((2022 - birthyr) %in% 25:34, "25-34",
      #   ifelse((2022 - birthyr) %in% 35:49, "35-49",
      #   ifelse((2022 - birthyr) %in% 50:64, "50-64",
      #   ifelse((2022 - birthyr) %in% 65:200, "65+", "missing")))))),
      race = ifelse(race %in% c("Asian", "White", "Black",
        "Native American"), as.character(race), "All Other"),
      hispanic = ifelse(hispanic == "Yes" | race == "Hispanic", "Yes", "No"),
      race_ethnicity = ifelse(race %in% c("Asian", "White", "Black", "Hispanic",
        "Native American"), as.character(race), "All Other"),
      gender = ifelse(gender4 == "Woman", "Woman",
        ifelse(gender4 == "Man", "Man", "other gender")),
      education =
        ifelse(educ %in% c("No HS", "High school graduate"), "HS or Less",
        ifelse(educ %in% c("Some college", "2-year"), "Some College",
        ifelse(educ == "4-year", "College Grad",
        ifelse(educ =="Post-grad", "Post-Grad", "missing")))),
      region)]

  categories[, age_educ := paste(age_cat, education)]
  categories[, age_gender := paste(age_cat, gender)]
  categories[, age_race := paste(age_cat, race_ethnicity)]
  categories[, educ_gender := paste(education, gender)]
  categories[, educ_race := paste(education, race_ethnicity)]
  categories[, gender_race := paste(gender, race_ethnicity)]
  denom <- categories[, sum(weight)]
  age_margins <- categories[, sum(weight) / denom, age_cat][
    order(age_cat)]
  race_margins <- categories[, sum(weight) / denom, race][
    order(race)]
  race_ethnicity_margins <- categories[, sum(weight) / denom, race_ethnicity][
    order(race_ethnicity)]
  hispanic_margins <- categories[, sum(weight) / denom, hispanic][
    order(hispanic)]
  educ_margins <- categories[, sum(weight) / denom, education][
    order(education)]
  gender_margins <- categories[, sum(weight) / denom, gender][
    order(gender)]
  region_margins <- categories[, sum(weight) / denom, region][
    order(region)]
  # age_educ_margins <- categories[, sum(weight) / denom, age_educ][
  #   order(age_educ)]
  # age_gender_margins <- categories[, sum(weight) / denom, age_gender][
  #   order(age_gender)]
  # age_race_margins <- categories[, sum(weight) / denom, age_race][
  #   order(age_race)]
  # educ_gender_margins <- categories[, sum(weight) / denom, educ_gender][
  #   order(educ_gender)]
  # educ_race_margins <- categories[, sum(weight) / denom, educ_race][
  #   order(educ_race)]
  # gender_race_margins <- categories[, sum(weight) / denom, gender_race][
  #   order(gender_race)]
  setorder(region_margins, region)
  setorder(age_margins, age_cat)
  setorder(educ_margins, education)
  setorder(gender_margins, gender)
  setorder(race_margins, race)
  setorder(hispanic_margins, hispanic)
  setorder(race_ethnicity_margins, race_ethnicity)
  # setorder(age_educ_margins, age_educ)
  # setorder(age_gender_margins, age_gender)
  # setorder(age_race_margins, age_race)
  # setorder(educ_gender_margins, educ_gender)
  # setorder(educ_race_margins, educ_race)
  # setorder(gender_race_margins, gender_race)
  region_target <- region_margins$V1 / sum(region_margins$V1)
  age_cat_target <- age_margins$V1 / sum(age_margins$V1)
  education_target <- educ_margins$V1 / sum(educ_margins$V1)
  gender_target <- gender_margins$V1 / sum(gender_margins$V1)
  race_target <- race_margins$V1 / sum(race_margins$V1)
  hispanic_target <- hispanic_margins$V1 / sum(hispanic_margins$V1)
  race_ethnicity_target <- race_ethnicity_margins$V1 / sum(race_ethnicity_margins$V1)
  # age_educ_target <- age_educ_margins$V1 / sum(age_educ_margins$V1)
  # age_gender_target <- age_gender_margins$V1 / sum(age_gender_margins$V1)
  # age_race_target <- age_race_margins$V1 / sum(age_race_margins$V1)
  # educ_gender_target <- educ_gender_margins$V1 / sum(educ_gender_margins$V1)
  # educ_race_target <- educ_race_margins$V1 / sum(educ_race_margins$V1)
  # gender_race_target <- gender_race_margins$V1 / sum(gender_race_margins$V1)
  names(region_target) <- region_margins$region
  names(age_cat_target) <- age_margins$age_cat
  names(education_target) <- educ_margins$education
  names(gender_target) <- gender_margins$gender
  names(race_target) <- race_margins$race
  names(race_ethnicity_target) <- race_margins$race_ethnicity
  names(hispanic_target) <- hispanic_margins$hispanic
  # names(age_educ_target) <- age_educ_margins$age_educ
  # names(age_gender_target) <- age_gender_margins$age_gender
  # names(age_race_target) <- age_race_margins$age_race
  # names(educ_gender_target) <- educ_gender_margins$educ_gender
  # names(educ_race_target) <- educ_race_margins$educ_race
  # names(gender_race_target) <- gender_race_margins$gender_race
  targets <- list(
    region_target, age_cat_target, education_target, gender_target,
    race_target, hispanic_target, race_ethnicity_target
    # , age_educ_target, age_gender_target, age_race_target,
    # educ_gender_target, educ_race_target, gender_race_target
    )
  names(targets) <- c(
    "region", "age_cat", "education", "gender", "race", "hispanic",
    "race_ethnicity"#,
    # "age_educ", "age_gender", "age_race",
    # "educ_gender", "educ_race", "gender_race"
    )
  targets
}

# balance ----
equiv_t_test <- function(x, y, w_x, w_y, epsilon, alpha = .05)
{
  x <- x[!is.na(x)]
  y <- y[!is.na(y)]
  weights_x <- w_x[!is.na(x)] / sum(w_x[!is.na(x)])
  weights_y <- w_y[!is.na(y)] / sum(w_y[!is.na(y)])
  dbar <- wtd.mean(x, weights_x) - wtd.mean(y, weights_y)
  m <- as.double(length(x))
  n <- as.double(length(y))
  N <- m + n
  x_var <- wtd.var(x, weights = weights_x, normwt = TRUE)
  y_var <- wtd.var(y, weights = weights_y, normwt = TRUE)
  non_cent <- (m * n * epsilon ^ 2) / N
  critical_const <- sqrt(qf(alpha, 1, N - 2, non_cent))
  se <- sqrt((m - 1) * x_var + (n - 1) * y_var) / sqrt(m * n * (N - 2) / N)
  df <- N - 2
  t_stat <-  dbar / se
  p <- pf(abs(t_stat) ^ 2, 1, df, non_cent)
  pooled_sd <- sqrt((m - 1) * x_var + (n - 1) * y_var) / sqrt(N - 2)
  obs_smd <- dbar / pooled_sd
  inverted <- try(
    uniroot(function(x) {
      pf(
        abs(t_stat) ^ 2,
        1,
        N - 2,
        ncp = (m * n * x ^ 2) / N) -
        ifelse(
          pf(abs(t_stat) ^ 2, 1, N - 2, ncp = (m * n * 0 ^ 2) / N) < alpha,
          pf(abs(t_stat) ^ 2, 1, N - 2, ncp = (m * n * obs_smd ^ 2) / N),
          alpha)
    }, c(0, 10 * abs(t_stat)), tol = 0.0001)$root,
    silent = TRUE
  )
  if (class(inverted) == "try-error") {
    inverted = NA
  }
  rej <- abs(t_stat) <= critical_const
  data.table(
    inverted = inverted,
    inverted_scaled = inverted * pooled_sd,
    p = p,
    observed_smd = obs_smd, #res$obs_smd,
    observed_diff = dbar, #res$obs_diff,
    power = 2 * pt(critical_const, N - 2) - 1, #res$power,
    treated_sd = sqrt(x_var),
    control_sd = sqrt(y_var), # res$sd
    pooled_sd = pooled_sd,
    se = se,
    tol = epsilon,
    t_stat = t_stat#,
    # critical_const = critical_const
  )
}
run_equiv <- function(X, Tr, w, epsilon.method = "std_effect", Y = NULL,
  custom_epsilon = NULL, fdr_correct = FALSE)
{
  switch(epsilon.method,
    std_effect = {
      tol =
        abs(mean(Y[Tr == 1], na.rm = TRUE) - mean(Y[Tr == 0], na.rm = TRUE)) /
        sd(Y[Tr == 0], na.rm = TRUE)
    },
    custom = {
      if (is.null(custom_epsilon))
        stop(paste0("ERROR: Must enter a custom epsilon value to use the ",
          "'custom' epsilon.method."))
      tol = custom_epsilon
    },
    strict = {
      tol = 0.36
    },
    liberal = {
      tol = 0.74
    },
    stop("ERROR: 'epsilon.method' not set to a valid option")
  )
  ranges <- rep(tol, ncol(X))
  names(ranges) <- names(X)
  tests <- do.call(rbind, lapply(names(X), function(var) {
    equiv_t_test(X[Tr == 1, get(var)], X[Tr == 0, get(var)], w[Tr == 1],
      w[Tr == 0], ranges[var])
  }))
  tests$variable_name <- names(ranges)
  setcolorder(tests,
    c("variable_name", "inverted", "inverted_scaled",
      "p", "observed_smd", "observed_diff",
      "power", "treated_sd", "control_sd", "pooled_sd", "se", "t_stat", "tol"))
  if (fdr_correct) {
    tests[, fdr_p := p.adjust(p, method = "BH")]
  }
  return(tests)
}
generate_equivalence_test_plot <- function(equiv_tests,
  panel_widths = c(1, 1, 5, 1, 1), display_names = NULL, var_rounding = 2,
  pval_rounding = 4, fdr_correct = TRUE, title_text = "") {
  require(cowplot)
  .e <- environment()
  if (!is.null(display_names)) {
    equiv_tests$display_name = factor(display_names,
      levels = rev(display_names))
  } else {
    equiv_tests$display_name = factor(equiv_tests$variable_name,
      levels = rev(equiv_tests$variable_name))
  }
  equiv_tests$const = rep(1, nrow(equiv_tests))
  equiv_tests = equiv_tests[nrow(equiv_tests):1, ]

  g <- ggplot(equiv_tests, aes(x = display_name)) +
    geom_linerange(aes(ymin = -inverted, ymax = inverted),
      linewidth = 5, color = "darkgray", alpha = 0.9)
  if (length(unique(equiv_tests$tol)) > 1) {
    g <- g +
      geom_linerange(aes(ymin = -tol, ymax = tol), linewidth = 10,
        color = 'gray')
  } else {
    lines <- unique(equiv_tests$tol)
    g = g + geom_hline(yintercept = c(-lines, lines), linetype = 2,
      linewidth = 0.75)
  }
  g <- g + scale_shape_identity() +
    geom_point(aes(y = observed_smd), color = "black", shape = 18, size = 4) +
    theme_bw() +
    coord_flip() +
    theme(
      axis.text.y = element_blank(),
      axis.ticks.y = element_blank(),
      axis.title.y = element_blank(),
      axis.text.x = element_text(size = 10),
      axis.title.x = element_text(size = 10),
      plot.title = element_text(size = 12, hjust = .5)) +
    labs(
      x = NULL,
      y = "Standard Deviations",
      font = 5)
  if (length(unique(equiv_tests$tol)) == 1) {
    g <- g + ggtitle(paste0("  \n", "\n",
      "Equivalence Range: +/- ", round(unique(equiv_tests$tol), 2),
      "\u03C3 "))
  } else {
    g <- g + ggtitle(paste0("\n \n \nEquivalence Tests"))
  }
  g_inv <- ggplot(equiv_tests,
    aes(x = display_name, y = const,
      label = sprintf(paste0("%0.0", var_rounding, "f"),
        round(inverted_scaled, var_rounding))),
    environment = .e) +
    ggtitle("Equivalence\nConfidence\nInterval (+/-)") + #\n(Scale of Var)") +
    geom_text() +
    theme_bw() +
    coord_flip() +
    theme(
      panel.grid.minor = element_blank(),
      panel.grid.major=element_blank(),
      axis.text.x = element_text(color = "white", size = 10),
      axis.ticks.x = element_line(color = "white"),
      axis.text.y = element_blank(),
      axis.ticks.y = element_blank(),
      axis.title.x = element_text(size = 10),
      plot.title = element_text(size = 12, hjust = .5)
    ) +
    ylim(1 - .05, 1.05) +
    labs(y = " ", x = NULL)
  g_obs <- ggplot(equiv_tests,
    aes(x = display_name, y = const,
      label = sprintf(paste0("%0.0", var_rounding, "f"),
        round(observed_diff, var_rounding))),
    environment = .e) +
    ggtitle("Observed\nMean\nDifference") + #\n(Scale of Var)") +
    geom_text() +
    theme_bw() +
    coord_flip() +
    theme(
      panel.grid.minor = element_blank(),
      panel.grid.major=element_blank(),
      axis.text.x = element_text(color = "white", size = 10),
      axis.ticks.x = element_line(color = "white"),
      axis.text.y = element_blank(),
      axis.ticks.y = element_blank(),
      axis.title.x = element_text(size = 10),
      plot.title = element_text(size = 12, hjust = .5)
    ) +
    ylim(1 - .05, 1.05) +
    labs(y = " ", x = NULL)
  fmt <- paste0("%.0", pval_rounding, "f")
  g_pval <- ggplot(equiv_tests,
    aes(x = display_name, y = const,
      label = ifelse(
        round(p, pval_rounding) == 0,
        paste0("< ", sprintf(fmt, 10 ^ -pval_rounding)),
        paste0("   ", sprintf(fmt, round(p, pval_rounding))))),
    environment = .e) +
    ggtitle(ifelse(fdr_correct, "FDR\nCorrected\nP-value", "\n\nP-value")) +
    geom_text() +
    theme_bw() +
    coord_flip() +
    theme(
      panel.grid.minor = element_blank(),
      panel.grid.major=element_blank(),
      axis.text.x = element_text(color = "white", size = 10),
      axis.ticks.x = element_line(color = "white"),
      axis.text.y = element_blank(),
      axis.ticks.y = element_blank(),
      axis.title.x = element_text(size = 10),
      plot.title = element_text(size = 12, hjust = .5)
    ) +
    ylim(1 - .05, 1.05) +
    labs(y = " ", x = NULL)
  g_var <- ggplot(equiv_tests,
    aes(x = display_name, y = const - .05, label = display_name)) +
    ggtitle(" \n \n   Variable") +
    geom_text(hjust = 0) +
    theme_bw() +
    coord_flip() +
    theme(
      panel.grid.minor = element_blank(),
      panel.border = element_blank(),
      panel.grid.major = element_blank(),
      axis.text.x = element_text(color = "white", size = 10),
      axis.ticks.x = element_line(color = "white"),
      axis.text.y = element_blank(),
      axis.ticks.y = element_blank(),
      axis.title.x = element_text(size = 10),
      plot.title = element_text(size = 12)
    ) +
    ylim(1 - .05, 1.05) +
    labs(y = " ", x = NULL)
  cs <- c(0, cumsum(panel_widths) / sum(panel_widths))
  x <- head(cs, -1)
  w <- tail(cs, -1) - head(cs, -1)
  h <- .9
  ggdraw() +
    draw_plot(g_var,  x = x[1], y = 0, height = h, width = w[1]) +
    draw_plot(g_obs,  x = x[2], y = 0, height = h, width = w[2]) +
    draw_plot(g,      x = x[3], y = 0, height = h, width = w[3]) +
    draw_plot(g_inv,  x = x[4], y = 0, height = h, width = w[4]) +
    draw_plot(g_pval, x = x[5], y = 0, height = h, width = w[5]) +
    draw_plot_label(x = .5, y = (h + 1) / 2, title_text, hjust = .5)
}
